package com.jboltai.capability.model.baichuan;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.message.AIMessageRole;
import com.jboltai.capability.message.MessageFinishReason;
import com.jboltai.capability.model.AIEventProcessor;
import com.jboltai.config.AIParamKey;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.resource.ai.AIResourceCenter;
import com.jboltai.util.http.HttpUtil;
import com.jboltai.util.http.RequestInstance;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;

public class BaichuanAIImpl implements AIEventProcessor {
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    private static final String STOP_FLAG = "[DONE]";

    public static final String API_URL = "https://api.baichuan-ai.com/v1/chat/completions";

    public static final BaichuanAIImpl INSTANCE = new BaichuanAIImpl();


    private RequestInstance buildRequest(AIEvent event, AIResource resource) {
        //资源指定了apiurl就用资源的，否则用默认的
        String apiUrl = StrUtil.isNotBlank(resource.getApiUrl()) ? resource.getApiUrl() : API_URL;
        RequestInstance request = HttpUtil.post(resource.getClient(), apiUrl, event);
        request.setAuthorization(resource.getApiKey());
        request.addJsonParam(AIParamKey.MODEL, event.getModelName());
        if (event.getTemperature() != null) {
            request.addJsonParam(AIParamKey.TEMPERATURE, event.getTemperature());
        }
        if (event.getTopP() != null) {
            request.addJsonParam(AIParamKey.TOP_P, event.getTopP());
        }
        if (event.getTopK() != null) {
            request.addJsonParam(AIParamKey.TOP_K, event.getTopK());
        }

        //处理messages
        List<AIMessage> messages = event.getMessages();
        if (messages != null && !messages.isEmpty()) {
            JSONArray messageArray = new JSONArray(messages.size());
            //百川不支持system 角色，要特殊转换一下
            processSystemMessage(event.getModelName(), messages);
            for (AIMessage message : messages) {
                messageArray.add(message.toJSONObject());
            }
            request.addJsonParam(AIParamKey.MESSAGES, messageArray);
        }
        //其他参数
        Map<String, Object> params = event.getParams();
        if (params != null && !params.isEmpty()) {
            params.keySet().forEach(key -> {
                request.addJsonParam(key, params.get(key));
            });
        }
        return request;
    }

    /**
     * 处理system 角色消息
     *
     * @param modelName 大模型名称
     * @param messages
     */
    private void processSystemMessage(String modelName, List<AIMessage> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }
        AIMessage systemMessage = messages.get(0);
        if (
                systemMessage != null
                && systemMessage.getRole() == AIMessageRole.system
                && (modelName.startsWith("Baichuan2") || modelName.startsWith("Baichuan3") )
        ) {
            //只有百川2不支持system消息，将system 拼到最后一条消息上
            messages.remove(0);
            AIMessage aiMessage = messages.get(messages.size() - 1);
            aiMessage.setContent(systemMessage.getContent() + ";" + aiMessage.getContent());
        }
    }

    @Override
    public void process(AIChatEvent event, AIResource resource) {


        RequestInstance request = buildRequest(event, resource);
        if (event.getMaxTokens() != null) {
            request.addJsonParam(AIParamKey.MAX_TOKENS, event.getMaxTokens());
        }

        event.setAttr(EventAttrKey.REQUEST_INSTANCE, request);

        LOGGER.debug("开始 调用Baichuan大模型 处理AIChat事件：\n\n\nEvent参数: {}, \n\n\nRequest参数: {}", JSON.toJSONString(event), JSON.toJSONString(request));
        boolean stream = event.getIsStream();
        if (stream) {
            //流式请求
            processEventByStream(event, resource, request);
        } else {
            processEventByNormal(event, resource, request);
        }
    }

    /**
     * 通过正常请求 处理event
     *
     * @param event
     * @param request
     */
    private static void processEventByNormal(AIChatEvent event, AIResource resource, RequestInstance request) {
        request.addJsonParam(AIParamKey.STREAM, false);
        request.asyncSend(response -> {
            try {
                String content = response.body().string();
                LOGGER.debug("AIChatEvent[{}]请求结束:{}", event.getId(), content);
                if (response.code() != 200) {
                    //有异常情况
                    handleFail(event, resource, response.code(), content, null);
                } else {
                    AIMessage message = handleChatSuccess(event, content, "message");
                    event.setResult(message);
                    event.setState(EventState.SUCCESS);

                }
            } catch (Exception e) {
                LOGGER.error("AIChatEvent[{}]处理响应结果发生异常:{}", event.getId(), e);
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                event.complete();
            }
        }, e -> {
            LOGGER.error("AIChatEvent[{}]请求失败:{}", event.getId(), e);
            event.fail(EventErrorType.NET_ERROR, e);
        });
    }

    /**
     * 以stream方式处理事件
     *
     * @param event
     * @param resource
     * @param request
     */
    private static void processEventByStream(AIChatEvent event, AIResource resource, RequestInstance request) {
        request.addJsonParam(AIParamKey.STREAM, true);
        request.asyncSendWithSSE(body -> {
            if (event.getState() != EventState.PROCESSING) {
                //中间发生异常，状态改变，不再处理响应
                return;
            }
            try {
                LOGGER.debug("AIChatEvent[{}]请求成功:{}", event.getId(), body);
                if (STOP_FLAG.equals(body)) {
                    event.setState(EventState.SUCCESS);
                    //结束
                    return;
                }
                handleChatSuccess(event, body, "delta");
            } catch (Exception e) {
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
                //todo 这里应该中断请求
            } finally {
                if (STOP_FLAG.equals(body)) {
                    event.complete();
                }
            }


        }, (res, e) -> {
            try {
                String response = null;
                if (res != null && res.body() != null) {
                    response = res.body().string();
                }
                handleFail(event, resource, res != null ? res.code() : 0, response, e);
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        });
    }

    /**
     * stream请求时，返回的内容放在"delta" 字段中，非stream请求，返回的内容放在message 字段中
     *
     * @param event
     * @param responseContent
     * @param contentField
     */
    private static AIMessage handleChatSuccess(AIChatEvent event, String responseContent, String contentField) {
        JSONObject data = JSON.parseObject(responseContent);
        AIMessage message = new AIMessage();
        message.setId(data.getString("id"));
        event.setResultId(message.getId());
        //获取回复
        JSONObject choice = data.getJSONArray("choices").getJSONObject(0);

        message.setContent(choice.getJSONObject(contentField).getString("content"));
        event.appendResponse(message.getContent());
        message.setRole(AIMessageRole.assistant);
        event.setResultRole(message.getRole());
        String finishReason = choice.getString("finish_reason");
        if (finishReason != null) {
            switch (finishReason) {
                case "stop":
                    event.setFinishReason(MessageFinishReason.STOP);
                    break;
                case "length":
                    event.setFinishReason(MessageFinishReason.LENGTH_LIMIT);
                    break;
                default:
                    event.setFinishReason(MessageFinishReason.OTHER);
            }
        }
        //获取使用情况
        data = data.getJSONObject("usage");
        if (data != null) {
            event.setPromptTokenCount(data.getInteger("prompt_tokens"));
            event.setTotalTokenCount(data.getInteger("total_tokens"));
            event.setCompletionTokenCount(data.getInteger("completion_tokens"));
        }

        event.success(message);
        return message;
    }

    /**
     * 失败时的处理
     *
     * @param event
     * @param resource
     * @param code
     * @param responseContent
     * @param exception
     */
    private static void handleFail(AIEvent event, AIResource resource, int code, String responseContent, Throwable exception) {
        JSONObject error = null;
        ;
        LOGGER.error("AIEvent[{}]请求失败:{}_{},异常：{}", event.getId(), code, responseContent, exception);
        if (exception == null) {
            //ai服务异常
            switch (code) {
                case 401:
                    event.fail(EventErrorType.API_KEY_INVALID, responseContent);
                    break;
                case 429: {
                    error = JSON.parseObject(responseContent);
                    if ("rate_limit".equals(error.getJSONObject("error").getString("code"))) {
                        AIResourceCenter.moveResourceToInvalid(resource);
                        event.fail(EventErrorType.EXCEED_REQUEST_LIMIT, responseContent);
                    } else {
                        event.fail(EventErrorType.ACCOUNT_BALANCE_INSUFFICIENT, responseContent);
                    }
                    break;
                }
                case 500:
                    event.fail(EventErrorType.SERVICE_SERVER_ERROR);
                    break;
                default:
                    //未知错误
                    event.fail(EventErrorType.UNKNOWN, responseContent);
                    break;
            }
        } else {
            //程序异常
            event.fail(EventErrorType.NET_ERROR, exception);
        }
    }


    @Override
    public void handleCancel(AIChatEvent event, AIResource resource) {
        LOGGER.warn("{}事件取消执行", event.getId());
        RequestInstance requestInstance = event.getAttr(EventAttrKey.REQUEST_INSTANCE);
        requestInstance.cancal();
    }
}
